import pandas as pd
import jieba
import re
import json
from rank_bm25 import BM25Okapi
from collections import Counter, defaultdict
from utils.utils import get_table_desc,show_all_tables
from Seacher.Query import search_columns
from utils.utils import grouped_dict
from math import log2
import jieba.analyse
def split_text(text):
    tokens = re.findall('[\u4e00-\u9fa5a-zA-Z]', text)
    tokens.extend(jieba.analyse.textrank(text)[:3])
    return tokens

class DocSearcher:
    def __init__(self, file_path, query_sample_fuc=None):
        self.table_name2db_name = {}
        self.table_name2desc = {}
        self.grouped_dict = grouped_dict
        self.search_for_bm25_token = []
        self.search_for_bm25_table = []
        self.search_for_bm25_column = []
        self.bm25 = None
        self.search_result = None
        self.return_inx = 0
        self.load_data(file_path)
        self.query_sample_fuc = query_sample_fuc
        # self.score_weight = {}
    def load_data(self, file_path):
        # Load the main data dictionary
        data = pd.read_excel(file_path)
        self.table_name2db_name = dict(zip(data['表英文'], data['库名英文']))
        data = data.to_dict(orient='records')

        # Create table descriptions
        for record in data:
            description = '\n'.join(f"{key}:{value}" for key, value in record.items())
            self.table_name2desc[record['表英文']] = description

        # Load column information and group by table name
        table2column = pd.read_excel(file_path, sheet_name='表字段信息')
        table2column.drop_duplicates(subset=['table_name','column_name'])
        self.score_weight = dict(Counter(table2column['column_name']))
        table2column['column_name'] = table2column['column_name'].fillna('无')

        # Prepare data for BM25 search
        for table_name, desc in self.table_name2desc.items():
            self._add_to_search_index(table_name, desc, 'table')

            for col_info in self.grouped_dict.get(table_name, []):
                col_desc = ' '.join(
                    re.sub(r'(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])', ' ', v).strip() if re.search('^[a-zA-Z]+$',
                                                                                                         v) else
                    ' '.join(split_text(v))
                    for v in col_info.values() if isinstance(v, str)
                )
                self._add_to_search_index(table_name, col_desc, col_info.get('column_name', '无'))

        # Initialize BM25 model
        self.bm25 = BM25Okapi(self.search_for_bm25_token)

    def _add_to_search_index(self, table_name, text, column_name):
        tokenized_text = [i for i in text.split(' ') if i]
        self.search_for_bm25_token.append(tokenized_text)
        self.search_for_bm25_table.append(table_name)
        self.search_for_bm25_column.append(column_name)

    def search_fields(self, query_tokens, top_n=5):
        # query_tokens = split_text(query)
        doc_scores = self.bm25.get_scores(query_tokens)
        top_n_indices = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True)[:top_n]
        return top_n_indices, [doc_scores[i] for i in top_n_indices]

    def search_by_query(self, query, top_n=20, score_threshold=5):
        query_token = split_text(query)
        top_n_indices, doc_scores = self.search_fields(query_token, top_n=top_n)

        recall_table_column = defaultdict(list)
        table_count_list = []

        for index, score in zip(top_n_indices, doc_scores):
            if score > score_threshold:
                table_count_list.append(self.search_for_bm25_table[index])
                recall_table_column[self.search_for_bm25_table[index]].append(self.search_for_bm25_column[index])

        results = []
        for table, cols in recall_table_column.items():
            results.append(get_table_desc(table, list(set(cols))))

        # Reset return_inx when a new search is performed
        self.return_inx = 0
        self.search_result = results
        return self._get_results_slice(3)




    @staticmethod
    def judge_tab_pos( market, table):
        if market == '':
            return True
        market_map = {'A股': ['HK_SecuMain', 'HK_EmployeeChange', 'HK_StockArchives', 'CS_HKStockPerformance',
                              'US_CompanyInfo', 'US_DailyQuote', 'US_SecuMain'],
                      '港股': ['SecuMain', 'LC_StockArchives', 'QT_StockPerformance', 'QT_DailyQuote',
                               'US_CompanyInfo', 'US_DailyQuote', 'US_SecuMain'
                               ],
                      '美股': [
                          'SecuMain', 'LC_StockArchives', 'QT_StockPerformance', 'QT_DailyQuote',
                          'HK_SecuMain', 'HK_EmployeeChange', 'HK_StockArchives', 'CS_HKStockPerformance',
                      ]}
        if table in market_map[market]:
            return False

        return True

    def search_by_field(self, query, field_top_n=3, score_threshold=5):

        if query[2] == '-':
            market, query = query.split('-', maxsplit=1)
        else:
            market = ''
        recall_table_column = defaultdict(list)
        # query_list = get_query(query, self.query_sample_fuc(query, 7))
        query_list = search_columns(query)
        table_score = defaultdict(float)

        for i in query_list:
            # if i['type'] == 'results':
            for tb, column in zip(self.search_for_bm25_table, self.search_for_bm25_column):

                if i['column_name_en'] == column:

                    score = 77/log2(self.score_weight.get(column,1)+1)
                    recall_table_column[tb].append({"column_name":column, "score":score})
                    table_score[tb] += score

            query_token = split_text(i['column_name_zh'])
            top_n_indices, doc_scores = self.search_fields(query_token, top_n=field_top_n)
            table_count_list = []
            for index, score in zip(top_n_indices, doc_scores):
                if score > score_threshold:
                    table_count_list.append(self.search_for_bm25_table[index])
                    recall_table_column[self.search_for_bm25_table[index]].append({
                        'column_name':self.search_for_bm25_column[index],
                        'score':score
                    })
                    table_score[self.search_for_bm25_table[index]] += score

        # print(table_score)

        results = []
        sorted_keys = sorted(table_score, key=table_score.get, reverse=True)
        for key in sorted_keys:
            if self.judge_tab_pos(market, key):
                column_recall = recall_table_column[key]
                sorted_column_recall = sorted(column_recall, key=lambda x: x['score'], reverse=True)
                column_lst = [i['column_name'] for i in sorted_column_recall]
                results.append(get_table_desc(key, column_lst))
            else:
                print(key,'不在',market)
        # Reset return_inx when a new search is performed
        self.return_inx = 0
        self.search_result = results
        return self._get_results_slice(3)

    def search_by_question(self, query, method='field'):
        if method == 'field':
            return self.search_by_field(query)
        elif method == 'query':
            return self.search_by_query(query)
        else:
            raise NotImplementedError

    def next_page(self, page_size=3):
        start_inx = self.return_inx
        end_inx = start_inx + page_size
        self.return_inx = end_inx
        return self._get_results_slice(page_size, start_inx) + '''\n如果此页依旧没有合适的文档，你可以使用：
```python
doc_searcher.search_by_query("想要搜索的字段中文说明")
```
搜索一个字段。
或者使用:
```python
doc_searcher.show_all_tables()
```
查看所有的表格,然后使用：
```python
doc_searcher.show_table("表英文")
```
来查看表格字段。
'''

    def _get_results_slice(self, page_size, start_inx=None):
        if start_inx is None:
            start_inx = self.return_inx
        res = '\n---\n'.join(self.search_result[start_inx:start_inx + page_size]) if self.search_result else ''

        return f'数据库文档:\n{res}\n'

    def choose(self, table_list):
        res = []
        try:
            for i in table_list:
                res.append(get_table_desc(i['table_name'], i['columns']))

            return '\n'.join(res)
        except Exception as e:
            return f'error\n{e}\n【重要提示】：你选择的表格有误，请重新选择或者搜索'

    @staticmethod
    def show_table(table_name):
        try:
            return get_table_desc(table_name)
        except KeyError:
            return '表格不存在，请使用`show_all_tables`查看所有表格。'

    @staticmethod
    def show_all_tables():
        return show_all_tables()


